"""
Nonlinear Helmholtz Equation solver using locELM.

This script reproduces the experiments from the paper for the 1D nonlinear Helmholtz equation:
    d^2u/dx^2 - lambda*u + beta*sin(u) = f(x)

Domain: x in [0, 8]
Parameters: lambda = 50, beta = 10
Exact solution: u(x) = sin(3*pi*x + 3*pi/20)*cos(4*pi*x - 2*pi/5) + 1.5 + x/10
"""

import numpy as np
import matplotlib.pyplot as plt
import time
import sys
import os

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from locELM.core.domain import Domain
from locELM.core.networks import MultiSubdomainNetwork
from locELM.solvers.nonlinear import NonlinearPDESolver
from locELM.utils.helpers import create_evaluation_function, compute_errors, Timer


class NonlinearHelmholtz1DSolver:
    """Solver for 1D nonlinear Helmholtz equation using locELM."""

    def __init__(self, lam=50.0, beta=10.0, a=0.0, b=8.0, seed=12):
        """Initialize solver parameters."""
        self.lam = lam
        self.beta = beta
        self.a = a
        self.b = b
        self.seed = seed

    def exact_solution(self, x):
        """Compute exact solution."""
        if isinstance(x, (int, float)):
            x = np.array([x])
        return (np.sin(3*np.pi*x + 3*np.pi/20) *
                np.cos(4*np.pi*x - 2*np.pi/5) + 1.5 + x/10)

    def source_term(self, x):
        """Compute source term f(x) such that exact solution satisfies PDE."""
        if isinstance(x, (int, float)):
            x = np.array([x])

        u = self.exact_solution(x)

        # Compute d^2u/dx^2 analytically
        d2u_dx2 = (
            -(3*np.pi)**2 * np.sin(3*np.pi*x + 3*np.pi/20) * np.cos(4*np.pi*x - 2*np.pi/5)
            - 2 * 3*np.pi * 4*np.pi * np.cos(3*np.pi*x + 3*np.pi/20) * np.sin(4*np.pi*x - 2*np.pi/5)
            - (4*np.pi)**2 * np.sin(3*np.pi*x + 3*np.pi/20) * np.cos(4*np.pi*x - 2*np.pi/5)
        )

        # f = d^2u/dx^2 - lambda*u + beta*sin(u)
        f = d2u_dx2 - self.lam * u + self.beta * np.sin(u)
        return f

    def solve(self, n_subdomains=4, n_colloc=100, M=200, R_m=5.0,
              method='nlsq_perturb', max_iter=50):
        """
        Solve the nonlinear Helmholtz equation.

        Parameters
        ----------
        n_subdomains : int
            Number of subdomains
        n_colloc : int
            Number of collocation points per subdomain
        M : int
            Number of training parameters (hidden layer width)
        R_m : float
            Maximum magnitude of random coefficients
        method : str
            Nonlinear solver method ('nlsq_perturb' or 'newton_llsq')
        max_iter : int
            Maximum iterations for nonlinear solver

        Returns
        -------
        dict
            Results including errors, timing, and networks
        """
        print(f"\n{'='*70}")
        print(f"1D Nonlinear Helmholtz Equation - locELM Method")
        print(f"{'='*70}")
        print(f"Domain: [{self.a}, {self.b}]")
        print(f"Number of subdomains: {n_subdomains}")
        print(f"Collocation points per subdomain: {n_colloc}")
        print(f"Training parameters per subdomain: {M}")
        print(f"Random coefficient magnitude (Rm): {R_m}")
        print(f"Nonlinear solver: {method}")
        print(f"{'='*70}\n")

        start_time = time.time()

        # Create domain
        domain = Domain(bounds=[(self.a, self.b)], n_subdomains=[n_subdomains])

        # Create networks
        networks = MultiSubdomainNetwork(
            n_subdomains=n_subdomains,
            input_dim=1,
            output_dim=1,
            hidden_layers=[M],
            activation='tanh',
            r_max=R_m,
            seed=self.seed
        )

        # Add networks for each subdomain
        for idx in range(n_subdomains):
            region = domain.get_subdomain_region(idx)
            networks.add_subdomain_network(region)

        # Create solver
        solver = NonlinearPDESolver(
            domain=domain,
            networks=networks,
            continuity_order=[1],  # C^1 continuity for 2nd order PDE
            collocation_type='uniform',
            n_collocation_points=[n_colloc]
        )

        # Define nonlinear operator
        def nonlinear_operator(u_values):
            """Nonlinear term: -lambda*u + beta*sin(u)"""
            return -self.lam * u_values + self.beta * np.sin(u_values)

        # Define nonlinear operator derivative
        def nonlinear_operator_derivative(u_values):
            """Derivative of nonlinear term: -lambda + beta*cos(u)"""
            return -self.lam + self.beta * np.cos(u_values)

        # Define boundary conditions
        boundary_conditions = {
            0: {  # x-dimension
                'lower': lambda x: np.array([self.exact_solution(self.a)]),
                'upper': lambda x: np.array([self.exact_solution(self.b)])
            }
        }

        # Solve
        print(f"Solving nonlinear PDE using {method}...")
        with Timer() as timer:
            if method == 'nlsq_perturb':
                weights, iterations = solver.solve_nlsq_perturb(
                    second_derivative_op=True,
                    nonlinear_op=nonlinear_operator,
                    source_term=self.source_term,
                    boundary_conditions=boundary_conditions,
                    output_dim=1,
                    delta=0.2,
                    xi2=1.0,
                    max_iter=max_iter,
                    tol=1e-10
                )
            else:  # newton_llsq
                weights, iterations = solver.solve_newton_llsq(
                    second_derivative_op=True,
                    nonlinear_op=nonlinear_operator,
                    nonlinear_op_derivative=nonlinear_operator_derivative,
                    source_term=self.source_term,
                    boundary_conditions=boundary_conditions,
                    output_dim=1,
                    max_iter=max_iter,
                    tol=1e-10
                )

        train_time = timer.elapsed_time
        print(f"Training completed in {train_time:.4f} seconds ({iterations} iterations)")

        # Evaluate solution
        eval_func = create_evaluation_function(networks, domain, output_dim=1)

        # Compute errors
        errors = compute_errors(
            solution_func=eval_func,
            exact_func=self.exact_solution,
            domain_bounds=[(self.a, self.b)],
            n_eval_points=[1000]
        )

        print(f"\n{'='*70}")
        print(f"Results")
        print(f"{'='*70}")
        print(f"Maximum Error: {errors['max_error']:.6e}")
        print(f"RMS Error: {errors['rms_error']:.6e}")
        print(f"Training Time: {train_time:.4f} seconds")
        print(f"Iterations: {iterations}")
        print(f"{'='*70}\n")

        return {
            'errors': errors,
            'timing': train_time,
            'iterations': iterations,
            'weights': weights,
            'networks': networks,
            'domain': domain,
            'eval_func': eval_func
        }

    def plot_solution(self, results, save_path=None):
        """Plot solution and error."""
        eval_func = results['eval_func']

        x = np.linspace(self.a, self.b, 500)
        u_computed = np.array([eval_func(np.array([[xi]])) for xi in x])
        u_exact = self.exact_solution(x)
        error = np.abs(u_computed - u_exact)

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

        # Solution plot
        ax1.plot(x, u_computed, 'b-', linewidth=2, label='Computed')
        ax1.plot(x, u_exact, 'r--', linewidth=2, label='Exact')
        ax1.set_xlabel('x', fontsize=12)
        ax1.set_ylabel('u', fontsize=12)
        ax1.set_title('Nonlinear Helmholtz Solution', fontsize=14)
        ax1.legend()
        ax1.grid(True)

        # Error plot
        ax2.semilogy(x, error, 'g-', linewidth=2)
        ax2.set_xlabel('x', fontsize=12)
        ax2.set_ylabel('Absolute Error', fontsize=12)
        ax2.set_title('Error Distribution', fontsize=14)
        ax2.grid(True)

        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"Plot saved to {save_path}")
        else:
            plt.show()


if __name__ == '__main__':
    solver = NonlinearHelmholtz1DSolver(lam=50.0, beta=10.0, a=0.0, b=8.0, seed=12)

    # Test case from paper: 4 subdomains, 100 colloc points, 200 training params
    print("\n" + "="*70)
    print("Reproducing Paper Results - Nonlinear Helmholtz")
    print("="*70)

    # Test with NLSQ-perturb method
    results_nlsq = solver.solve(
        n_subdomains=4,
        n_colloc=100,
        M=200,
        R_m=5.0,
        method='nlsq_perturb',
        max_iter=100
    )

    solver.plot_solution(results_nlsq,
                         save_path='experiments/nonlinear_helmholtz_nlsq.png')

    # Test with Newton-LLSQ method
    results_newton = solver.solve(
        n_subdomains=4,
        n_colloc=100,
        M=200,
        R_m=5.0,
        method='newton_llsq',
        max_iter=100
    )

    solver.plot_solution(results_newton,
                         save_path='experiments/nonlinear_helmholtz_newton.png')

    print("\n" + "="*70)
    print("Method Comparison")
    print("="*70)
    print(f"NLSQ-perturb: Max Error = {results_nlsq['errors']['max_error']:.6e}, "
          f"Time = {results_nlsq['timing']:.4f}s, Iterations = {results_nlsq['iterations']}")
    print(f"Newton-LLSQ:  Max Error = {results_newton['errors']['max_error']:.6e}, "
          f"Time = {results_newton['timing']:.4f}s, Iterations = {results_newton['iterations']}")
    print("="*70)
